package RPC

import (
	pb "AttackTracer/RpcMessage/AttackTracer/RpcMessage"
	"AttackTracer/TracingEngine/config"
	"context"
	"errors"
	"google.golang.org/grpc"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/metadata"
	"log"
	"net"
	"time"
)

var (
	kaep = keepalive.EnforcementPolicy{
		MinTime:             5 * time.Second, // If a client pings more than once every 5 seconds, terminate the connection
		PermitWithoutStream: true,            // Allow pings even when there are no active streams
	}
	kasp = keepalive.ServerParameters{
		MaxConnectionIdle:     15 * time.Second, // If a client is idle for 15 seconds, send a GOAWAY
		MaxConnectionAge:      30 * time.Second, // If any connection is alive for more than 30 seconds, send a GOAWAY
		MaxConnectionAgeGrace: 5 * time.Second,  // Allow 5 seconds for pending RPCs to complete before forcibly closing connections
		Time:                  5 * time.Second,  // Ping the client if it is idle for 5 seconds to ensure the connection is still active
		Timeout:               1 * time.Second,  // Wait 1 second for the ping ack before assuming the connection is dead
	}
	rpcConfig     config.Rpc
	AnalyzerConns = make(map[string]time.Time)
)

type RpcServer struct {
	pb.UnimplementedAnalyzerTracingEngineServer
}

// SayHello implements AnalyzerTracingEngineServer
func (s *RpcServer) SayHello(ctx context.Context, in *pb.HelloRequest) (*pb.HelloReply, error) {
	log.Printf("Hello received from %v", in.GetId())
	return &pb.HelloReply{Message: "Hello from " + rpcConfig.Id + " to " + in.GetId()}, nil
}

func Serve(rpcConf config.Rpc) {
	rpcConfig = rpcConf
	lis, err := net.Listen("tcp", rpcConfig.ListenAddr)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	log.Println("RPC listening on", rpcConfig.ListenAddr, "as", rpcConfig.Id)
	s := grpc.NewServer(grpc.KeepaliveEnforcementPolicy(kaep),
		grpc.KeepaliveParams(kasp),
		grpc.StreamInterceptor(streamInterceptor), // credential handling for stream
		grpc.UnaryInterceptor(unaryInterceptor))   // credential handling for unary

	pb.RegisterAnalyzerTracingEngineServer(s, &RpcServer{})
	if err := s.Serve(lis); err != nil {
		log.Fatalf("failed to serve: %v", err)
	}
}

/********** login credentials handling **********/
func streamInterceptor(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	if err := authorize(stream.Context()); err != nil {
		return err
	}
	return handler(srv, stream)
}
func unaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
	if err := authorize(ctx); err != nil {
		return err, err
	}
	return handler(ctx, req)
}
func authorize(ctx context.Context) error {
	if md, ok := metadata.FromIncomingContext(ctx); ok {
		//log.Print(md)
		if len(md["accesskey"]) > 0 && md["accesskey"][0] == rpcConfig.SelfKey &&
			len(md["id"]) == 1 {
			AnalyzerConns[md["id"][0]] = time.Now() // add to records
			//log.Println(AnalyzerConns)
			return nil
		}
		return errors.New("AccessDenied")
	}

	return errors.New("EmptyMetadata")
}
